from pox.core import core
import pox.openflow.libopenflow_01 as of
from pox.lib.util import dpidToStr

log = core.getLogger()

def _handle_ConnectionUp(event):
    dpid = event.connection.dpid
    sw = dpidToStr(dpid)
    log.info("Switch %s has connected.", sw)
    if sw == '00-00-00-00-00-01':  # s1
        configure_switch(event.connection, is_s1=True)
    elif sw == '00-00-00-00-00-02':  # s2
        configure_switch(event.connection, is_s1=False)
    else:
        log.warning("Unknown switch connected: %s", sw)

def configure_switch(conn, is_s1):
    if is_s1:
        host_pairs = [
            ("192.168.0.101", "192.168.0.104", 1),
            ("192.168.0.102", "192.168.0.105", 2),
            ("192.168.0.103", "192.168.0.106", 3)
        ]
    else:
        host_pairs = [
            ("192.168.0.104", "192.168.0.101", 1),
            ("192.168.0.105", "192.168.0.102", 2),
            ("192.168.0.106", "192.168.0.103", 3)
        ]

    for src_ip, dst_ip, local_port in host_pairs:
        # 配置从源 IP 到目标 IP 的流表项
        msg = of.ofp_flow_mod()
        msg.priority = 1000
        msg.match.dl_type = 0x0800  # IPv4
        msg.match.nw_src = src_ip
        msg.match.nw_dst = dst_ip
        msg.actions.append(of.ofp_action_output(port=local_port))  # 使用local_port  
        try:
            conn.send(msg)
            log.info("Flow mod sent: src_ip=%s, dst_ip=%s, action_port=%d", src_ip, dst_ip, local_port)  
        except Exception as e:
            log.error("Failed to send flow mod: %s", e)

        # 配置从目标 IP 到源 IP 的流表项
        msg = of.ofp_flow_mod()
        msg.priority = 1000
        msg.match.dl_type = 0x0800  # IPv4
        msg.match.nw_src = dst_ip
        msg.match.nw_dst = src_ip
        msg.actions.append(of.ofp_action_output(port=local_port))  # 使用local_port  
        try:
            conn.send(msg)
            log.info("Flow mod sent: src_ip=%s, dst_ip=%s, action_port=%d", dst_ip, src_ip, local_port)  
        except Exception as e:
            log.error("Failed to send flow mod: %s", e)

    log.info("Configured IP-based flows for switch %s", "s1" if is_s1 else "s2")  

def launch():
    core.openflow.addListenerByName("ConnectionUp", _handle_ConnectionUp)
    log.info("POX IP-based communication module running.")